# 四足机器人小跑步态控制模块
from math import sin, cos, atan, pi
import time
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # 导入3D绘图模块
Ts = 1  # 步长时间

def cal_t(t, xs, xf, raiseH, r1, r2, r3, r4):    # 小跑步态执行函数
    """
    t: 时间参数，用于控制步态周期
    xs: 起始位置x坐标
    xf: 目标位置x坐标
    raiseH: 抬腿高度
    r1,r2,r3,r4: 方向控制参数，用于判断转向
    """
    Tm = Ts*0.5    # 步态周期时间
    Hc = 107    # 髋关节高度
    z = 0       # 转向角度
    S0 = xf     # 目标位置x坐标

    # 方向判断
    if r1 == -1 and r2 == -1 and r3 == 1 and r4 == 1:   # 左转
        z = -10
    elif r1 == 1 and r2 == 1 and r3 == -1 and r4 == -1:   # 右转
        z = 10

    H0 = raiseH  # 抬腿高度
    a = 0        # 初始角度
    d = atan(z / Hc)  # 转向角度计算 横向角度

    if t < Tm:  # 前半步态周期
        sigma = 2 * pi * t / Tm
        c = cos(d * cos(pi * t / Tm))  # 转向角度余弦值 角度越大补偿值越大 
        b = (1 - cos(sigma)) / 2  # 抬腿高度系数
        x0 = S0 * (t / Tm - 1 / (2 * pi) * sin(sigma)) - 0.5 * S0  # x轴位移
        y0 = (H0 * b) / c  # y轴位移 腿横向跨度越大，高度就越低，除以补偿值c就是弥补高度
        z0 = a + z * cos(pi * t / Tm)  # z轴位移
        x3 = x0
        y3 = y0
        z3 = z0
        x1 = -x0
        y1 = (-H0 / 40 * b) / c  # 支撑腿往下撑一点
        z1 = z0
        x2 = x1
        y2 = y1
        z2 = z0
        
    if t >= Tm:  # 后半步态周期
        sigma = 2 * pi * (t - Tm) / Tm
        c = cos(d * cos(pi * t / Tm))  # 转向角度余弦值
        b = (1 - cos(sigma)) / 2  # 抬腿高度系数
        x0 = 0.5 * S0 - S0 * ((t - Tm) / Tm - 1 / (2 * pi) * sin(sigma))  # x轴位
        y0 = (-H0 / 40 * b) / c  # y轴位移
        z0 = a + z * cos(pi * t / Tm)  # z轴位移
        x3 = x0
        y3 = y0
        z3 = z0
        x1 = -x0
        y1 = (H0 * b) / c
        z1 = z0
        x2 = x1
        y2 = y1
        z2 = z0
        
    # 返回所有关节的x,y,z坐标
    return x0, x1, x3, x2, y0, y1, y3, y2, z0, z1, z3, z2

# 初始化变量
t = 0

speed = 0.065  # 速度
results = []

L=1
R=-1
# 模拟数据
xs = 0
xf = 6*8
raiseH = 40
r1, r2, r3, r4 = L, L, R, R

def testloop():
    global t, Ts, speed, results
    if t < Ts:
        result = cal_t(t, xs, xf, raiseH, r1, r2, r3, r4)
        results.append(result)
        t += speed
        update_plot()
        plt.pause(0.03)  # 暂停以更新图形
        # 输出坐标点到控制台
        #print(f"Leg 0: ({result[0]}, {result[4]}, {result[8]})")
        #print(f"Leg 1: ({result[1]}, {result[5]}, {result[9]})")
        #print(f"Leg 2: ({result[2]}, {result[6]}, {result[10]})")
        #print(f"Leg 3: ({result[3]}, {result[7]}, {result[11]})")
    else:
        # 停止定时器
        plt.ioff()
        plt.show()
        
        return

# 设置交互模式
plt.ion()

# 创建图形
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')  # 创建3D子图

# 绘制坐标
def update_plot():
    global results
    if results:
        x0, x1, x2, x3, y0, y1, y2, y3, z0, z1, z2, z3 = zip(*results)
        ax.cla()
        # 交换 y 和 z 轴数据
        ax.plot(x0, z0, y0, label='Leg 0')
       # ax.plot(x1, z1, tuple(y+10 for y in y1), label='Leg 1')
        ax.plot(x2, z2, tuple(y-10 for y in y2), label='Leg 2')
       # ax.plot(x3, z3, y3, label='Leg 3')
        ax.set_title('3D Coordinates of Legs')
        ax.set_xlabel('X')
        ax.set_ylabel('Z')
        ax.set_zlabel('Y')
        ax.legend()
        ax.set_ylim(ax.get_ylim()[::-1])  # 反转 Y 轴方向，使其垂直向上

        fig.tight_layout()

# 主循环模式
while t < Ts:
    testloop()

input("按回车键退出...")  # 增加按回车键退出的功能

